import xgboost as xgb
from sklearn.datasets import load_diabetes
from sklearn.model_selection import train_test_split
from sklearn.metrics import root_mean_squared_error
import pandas as pd
import matplotlib.pyplot as plt

"""
XGBoost回归示例 - 使用糖尿病数据集
糖尿病数据集包含10个特征，目标变量是糖尿病进展的定量测量
特征数：10个（年龄、性别、BMI、血压等）
样本数：442个
目标变量：糖尿病进展的定量测量（连续值）
"""

# 使用糖尿病数据集（无需网络下载）
# load_diabetes() 函数加载内置的糖尿病数据集
# return_X_y=True 参数表示直接返回特征矩阵X和目标变量y
X, y = load_diabetes(return_X_y=True)  # y 是糖尿病进展的定量测量

# 展示数据类型
print("X的数据类型:", type(X))  # numpy.ndarray
print("y的数据类型:", type(y))  # numpy.ndarray
print("X的形状:", X.shape)
print("y的形状:", y.shape)
print("目标变量统计:")
print(f"  最小值: {y.min():.2f}")
print(f"  最大值: {y.max():.2f}")
print(f"  均值: {y.mean():.2f}")
print(f"  标准差: {y.std():.2f}")

# 划分数据集
# train_test_split() 函数将数据集划分为训练集和测试集
# test_size=0.3 表示测试集占总数据的30%，训练集占70%
# random_state=42 设置随机种子，确保每次划分结果一致，便于结果复现
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.3, random_state=42
)


def train_model(X_train, y_train, max_depth=3, n_estimators=100):
    """
    训练XGBoost回归模型
    :param X_train: 训练集特征
    :param y_train: 训练集目标值
    :param max_depth: 树的最大深度，默认值为3
    :param n_estimators: 弱学习器数量，默认值为100
    :return: 训练好的模型实例
    """

    # 创建XGBoost回归器实例
    # objective="reg:squarederror" 设置目标函数为平方误差回归
    reg = xgb.XGBRegressor(
        objective="reg:squarederror", max_depth=max_depth, n_estimators=n_estimators
    )
    # 使用训练数据拟合模型
    reg.fit(X_train, y_train)
    return reg

    # d_train = xgb.DMatrix(X_train, label=y_train)
    # params = {
    #     "objective": "reg:squarederror",
    #     "max_depth": max_depth,
    #     "n_estimators": n_estimators,
    # }
    # model = xgb.train(params, d_train)
    # return model


def evaluate_model(model, X_test, y_test):
    """
    评估模型性能
    :param model: 训练好的模型实例
    :param X_test: 测试集特征
    :param y_test: 测试集目标值
    :return: 均方根误差（RMSE）
    """
    # 使用测试集进行预测，并计算均方根误差（RMSE）
    rmse = root_mean_squared_error(y_test, model.predict(X_test))
    # RMSE是回归问题常用的评估指标，值越小表示模型预测越准确
    print("RMSE:", rmse)
    return rmse


def save_model(model, model_path):
    """
    保存模型到指定路径
    :param model: 训练好的模型实例
    :param model_path: 模型保存路径
    """
    if not model_path.endswith(".model"):
        model_path += ".model"
    model.save_model(model_path)
    print(f"模型已保存至 {model_path}")


def load_model(model_path):
    """
    从指定路径加载模型
    :param model_path: 模型保存路径
    :return: 加载的模型实例
    """
    model = xgb.XGBRegressor()
    model.load_model(model_path)
    print(f"模型已从 {model_path} 加载")
    return model


def predict(model, X):
    """
    使用模型进行预测
    :param model: 训练好的模型实例
    :param X: 输入特征数据
    :return: 预测结果
    """
    return model.predict(X)


def show_feature_importance(model):
    """
    显示特征重要性
    """
    diabetes = load_diabetes()
    feature_names = diabetes.feature_names

    # 获取特征重要性
    importance_scores = model.feature_importances_

    # 创建重要性表格
    importance_df = pd.DataFrame(
        {"特征": feature_names, "重要性": importance_scores}
    ).sort_values("重要性", ascending=False)

    print("\n=== 特征重要性排名 ===")
    for i, row in importance_df.iterrows():
        print(f"{i+1}. {row['特征']}: {row['重要性']:.4f}")

    # 简单可视化
    plt.figure(figsize=(10, 6))
    plt.barh(importance_df["特征"], importance_df["重要性"])
    plt.xlabel("Importance Score")
    plt.title("Diabetes Dataset Feature Importance")
    plt.tight_layout()
    plt.show()

    return importance_df


# 主程序执行
if __name__ == "__main__":
    # 训练模型
    print("开始训练模型...")
    model = train_model(X_train, y_train)

    # 评估模型
    print("模型评估结果:")
    evaluate_model(model, X_test, y_test)

    # 特征重要性分析
    print("\n进行特征重要性分析...")
    importance_df = show_feature_importance(model)

    # # 保存模型
    # save_model(model, "models/diabetes_xgboost_model.model")

    # # 加载模型
    # print("\n验证模型加载功能:")
    # loaded_model = load_model("models/diabetes_xgboost_model.model")

    # # 使用模型，预测测试集
    # y_pred = predict(model, X_test)
    # print("前5个测试样本的预测值:", y_pred[:5])
